package network

import (
	"crypto/tls"
	"github.com/gorilla/websocket"
	"leaf/core/log"
	"leaf/core/tool"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

type WsServer struct {
	addr       string
	NewAgent   func(conn Conn) Agent
	ln         net.Listener
	MaxConnNum int
	conns      map[*websocket.Conn]*WsConn
	mutexConns sync.Mutex
	wgLn       sync.WaitGroup
	wgConns    sync.WaitGroup
	OnAccept   func(conn Conn)
	msgParser  *MsgParser
	closeFlag  int32

	HTTPTimeout time.Duration
	CertFile    string
	KeyFile     string
	httpServer  *http.Server
	upgrader    websocket.Upgrader
}

func (server *WsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	conn, err := server.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Debug("upgrade error: %v", err)
		return
	}
	conn.SetReadLimit(int64(Default_MaxMsgLen))

	server.mutexConns.Lock()
	connsLen := len(server.conns)
	server.mutexConns.Unlock()
	if connsLen >= Default_MaxConnNum {
		conn.Close()
		log.Debug("too many connections")
		return
	}

	wsConn := newAcceptWsConn(conn, server.msgParser, server)
	if wsConn == nil {
		conn.Close()
		log.Debug("new ws accept conn fail")
		return
	}
	server.mutexConns.Lock()
	server.conns[conn] = wsConn
	server.mutexConns.Unlock()
	wsConn.agent = server.NewAgent(wsConn)
	server.wgConns.Add(1)
	if server.OnAccept != nil {
		server.OnAccept(wsConn)
	}
	wsConn.Start()
}

func NewWsServer(addr string, maxConnNum int, createAgent func(conn Conn) Agent) *WsServer {
	s := new(WsServer)
	s.addr = addr
	s.MaxConnNum = maxConnNum
	if s.MaxConnNum <= 0 {
		s.MaxConnNum = Default_MaxConnNum
	}
	if s.OnAccept == nil {

	}
	s.HTTPTimeout = 10 * time.Second
	s.conns = make(map[*websocket.Conn]*WsConn, s.MaxConnNum)
	s.NewAgent = createAgent
	s.upgrader = websocket.Upgrader{
		HandshakeTimeout: s.HTTPTimeout,
		CheckOrigin:      func(_ *http.Request) bool { return true },
	}
	return s
}

func (server *WsServer) Start(asyn bool) {
	ln, err := net.Listen("tcp", server.addr)
	if err != nil {
		log.Fatal("%v", err)
	}
	if server.NewAgent == nil {
		log.Fatal("NewAgent must not be nil")
	}
	if server.msgParser == nil {
		server.msgParser = DefaultMsgParser
	}
	if server.CertFile != "" || server.KeyFile != "" {
		config := &tls.Config{}
		config.NextProtos = []string{"http/1.1"}

		var err error
		config.Certificates = make([]tls.Certificate, 1)
		config.Certificates[0], err = tls.LoadX509KeyPair(server.CertFile, server.KeyFile)
		if err != nil {
			log.Fatal("%v", err)
		}
		ln = tls.NewListener(ln, config)
	}

	server.ln = ln
	server.httpServer = &http.Server{
		Addr:           server.addr,
		Handler:        server,
		ReadTimeout:    server.HTTPTimeout,
		WriteTimeout:   server.HTTPTimeout,
		MaxHeaderBytes: 1024,
	}
	log.Release("WsServer Start addr:%v", server.addr)
	accept_func := func() {
		server.wgLn.Add(1)
		defer func() {
			server.wgLn.Done()
			server.Close()
		}()

		server.httpServer.Serve(ln)
	}
	if !asyn {
		accept_func()
	} else {
		tool.Go(accept_func)
	}
}

func (server *WsServer) onCloseConn(wsConn *WsConn) {
	server.wgConns.Add(-1)
	server.mutexConns.Lock()
	delete(server.conns, wsConn.conn)
	server.mutexConns.Unlock()
	wsConn.Destroy()
}

func (server *WsServer) Close() {
	if !atomic.CompareAndSwapInt32(&server.closeFlag, 0, 1) {
		return
	}
	if server.ln != nil {
		server.ln.Close()
		server.ln = nil
	}
	//等待监听线程退出
	server.wgLn.Wait()

	var conns []*WsConn
	server.mutexConns.Lock()
	for _, wsConn := range server.conns {
		conns = append(conns, wsConn)
	}
	server.mutexConns.Unlock()

	//关闭所用连接
	for _, wsConn := range conns {
		wsConn.Close()
	}
	//等待所有连接线程返回
	server.wgConns.Wait()
	//销毁连接资源
	for _, wsConn := range conns {
		wsConn.Destroy()
	}
}
